#encoding=utf-8
import xml.etree.ElementTree as ET
import cv2
import os
import sys
import numpy as np
import glob

class_names = ('BACKGROUND',
                    '1111',
                    )
class_dict = {class_name: i for i, class_name in enumerate(class_names)}


def _get_annotation(annotation_file):
    et = ET.parse(annotation_file)
    objects = et.findall("object")
    size = et.findall("size")
    w = int(size[0].find('width').text)
    h = int(size[0].find('height').text)
    bboxs = []
    labels = []
    for object in objects:
        class_name = object.find('name').text.lower().strip()
        if class_name == '1111':
            bndbox = object.find('bndbox')
            # VOC dataset format follows Matlab, in which indexes start from 0

            xmin = int(bndbox.find('xmin').text)
            ymin = int(bndbox.find('ymin').text)
            xmax = int(bndbox.find('xmax').text)
            ymax = int(bndbox.find('ymax').text)

            bboxs.append([xmin, ymin, xmax, ymax])

            labels.append(class_dict[class_name])
    return bboxs, labels

if __name__ == "__main__":
    image_save_folder = '/media/handsome/backupdata/hanson/WIDER_train/images/62-Heils_Face'
    root_path = r'/media/handsome/backupdata/hanson/WIDER_train/others/人脸标完/人脸/need_label'
    anno_src_pth = os.path.join(root_path, "Annotations")
    img_src_pth = os.path.join(root_path, "Images")

    f = open('../anno_store/heils-face-annos.txt', 'w')
    for idx, folder in enumerate(os.listdir(anno_src_pth)):
        folder_pth = os.path.join(anno_src_pth,folder)
        print(folder_pth)
        for jdx, xml_path in enumerate(glob.glob(os.path.join(folder_pth, '*.xml'))):
            bboxs, labels = _get_annotation(xml_path)
            if len(bboxs) == 0:
                print('None box:%s'% xml_path)
                continue
            img_pth = os.path.join(img_src_pth,xml_path.split('/')[-2], xml_path.split('/')[-1].replace('.xml','.jpg'))
            img = cv2.imread(img_pth)
            img_name = image_save_folder.split('/')[-1] + '-' + img_pth.split('/')[-1]
            file_img_str = image_save_folder.split('/')[-1] + '/' + img_name
            f.write('%s '% file_img_str)
            for box in bboxs:

                # cv2.rectangle(img,(box[0], box[1]),(box[2],box[3]),(255,0,0),2)
                f.write('%d %d %d %d ' %
                        (box[0], box[1], box[2], box[3]))
            f.write('\n')
            save_image_path = os.path.join(image_save_folder, img_name)
            cv2.imwrite(save_image_path, img)
            # cv2.imshow('img',img)
            # cv2.waitKey(0)
